import numpy as np
from math import sin, cos, radians, sqrt
from typing import Tuple

# --- I. 物理常数和弹丸参数 ---
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 绝热垂直温度梯度 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 干燥空气比气体常数 (J/kg·K)
P_TERM = (G / (R * L)) - 1 # 用于计算密度比的指数
C_SOUND = 338.0      # 海平面声速 (m/s)

M = 5.0              # 弹丸质量 (kg)
A = 0.0095           # 弹丸截面积 (m^2)
V_WIND_X = 0.0       # 纵向风速度 (固定为零)

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) # 马赫数查找点
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) # 阻力系数

# --- II. 动态物理模型和RK4函数 (2D简化版) ---

def get_drag_coefficient(v_rel):
    """根据相对速度计算空气阻力系数 (Cd)。"""
    mach_number = v_rel / C_SOUND
    return np.interp(mach_number, MA_POINTS, CD_POINTS)

def get_air_density(z):
    """根据高度 Z 计算空气密度 (假设温度随高度变化)。"""
    h = max(0.0, z) 
    # 气压和温度随高度变化导致的密度变化
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0: return 0.0 
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

def derivatives_of_state_2d(S: np.ndarray, params: Tuple[float, float, float]) -> np.ndarray:
    """计算 2D 状态向量 S = [x, z, vx, vz] 在时间上的微分。"""
    x, z, vx, vz = S
    M, A, G = params 
    
    # 1. 计算相对速度 (纵向风 V_WIND_X = 0.0)
    vx_rel = vx - V_WIND_X 
    vz_rel = vz
    v_rel = sqrt(vx_rel**2 + vz_rel**2) # 相对速度的模

    if v_rel < 1e-6:
        return np.array([vx, vz, 0.0, -G])

    rho = get_air_density(z) # 密度随飞行高度 Z 变化
    Cd = get_drag_coefficient(v_rel)

    # 2. 空气阻力加速度
    K = -(Cd * rho * A * v_rel) / (2.0 * M)
    
    dvx_drag = K * vx_rel
    dvz_drag = K * vz_rel

    # 3. 合并加速度
    dvx_dt = dvx_drag
    dvz_dt = dvz_drag - G

    return np.array([vx, vz, dvx_dt, dvz_dt])

def rk4_step_2d(S: np.ndarray, dt: float, params: Tuple[float, float, float]) -> np.ndarray:
    """四阶龙格-库塔方法进行一步 2D 积分。"""
    f = lambda S: derivatives_of_state_2d(S, params)
    
    K1 = dt * f(S)
    K2 = dt * f(S + K1 / 2.0)
    K3 = dt * f(S + K2 / 2.0)
    K4 = dt * f(S + K3)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

def simulate_and_get_impact_2d(v0: float, theta_deg: float, dt: float, h0: float) -> float:
    """
    模拟 2D 轨迹并返回最终落点 X 坐标 (射程)。
    h0 参数定义了发射高度和目标平面的高度。
    """
    params = (M, A, G)
    theta_rad = radians(theta_deg)
    
    # 2D 状态向量: [x, z, vx, vz]. 初始高度 Z = H0
    S = np.array([
        0.0, h0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    # 假设炮弹在初始高度 H0 落地
    while S[1] >= h0 - 0.01:
        S_new = rk4_step_2d(S, dt, params) 
        
        # 冲击检查 (Z 穿过 H0 平面)
        if S_new[1] < h0 and S[1] >= h0:
            # 线性插值计算冲击点 X 坐标 (射程)
            # t_impact_ratio = (S[1] - h0) / (S[1] - S_new[1])
            t_impact_ratio = (S[1] - h0) / (S[1] - S_new[1])
            R_final = S[0] + (S_new[0] - S[0]) * t_impact_ratio
            return R_final
            
        S = S_new
        
    return S[0]

# --- III. 搜索函数 (计算 V0 和 theta) ---

def binary_search_v0_for_range(target_range: float, theta_deg: float, dt: float, h0: float, V_low=10.0, V_high=2000.0, tolerance=1.0) -> float:
    """查找固定角度、目标距离和初始高度所需的 V0。"""
    
    for _ in range(50):
        V_mid = (V_low + V_high) / 2.0
        # 模拟计算该 V_mid 在当前高度下的射程
        R_mid = simulate_and_get_impact_2d(V_mid, theta_deg, dt, h0)
        
        if abs(R_mid - target_range) < tolerance:
            return V_mid
        
        if R_mid < target_range:
            V_low = V_mid 
        else:
            V_high = V_mid 
            
    # 如果搜索失败，返回一个指示失败的值
    return np.nan

def search_min_v0_and_theta(target_range: float, dt: float, h0: float) -> Tuple[float, float]:
    """
    搜索达到目标射程所需的最小初速度 V0 及其对应的发射角 theta (在高度 h0)。
    """
    
    # 角度搜索范围和步长
    THETA_MIN = 20.0
    THETA_MAX = 70.0
    THETA_STEP = 0.5  # 角度步长

    best_v0 = float('inf')
    best_theta = np.nan
    
    # 遍历可能的发射角
    for theta_deg in np.arange(THETA_MIN, THETA_MAX + THETA_STEP, THETA_STEP):
        # 搜索当前角度所需的 V0
        v0_required = binary_search_v0_for_range(
            target_range=target_range, 
            theta_deg=theta_deg, 
            dt=dt, 
            h0=h0
        )
        
        # 检查是否找到有效的 V0 并且比当前最优 V0 小
        if not np.isnan(v0_required) and v0_required < best_v0:
            best_v0 = v0_required
            best_theta = theta_deg
                
    if best_v0 == float('inf'):
        return np.nan, np.nan
        
    return best_v0, best_theta

# --- IV. 主运行和表格输出 ---

def generate_altitude_table():
    
    # --- 固定参数 ---
    DT = 0.1             # RK4 步长 (s)

    # --- 表格参数 ---
    R_MIN = 1000.0
    R_MAX = 1500.0
    R_STEP = 100.0        # 目标距离步长
    H_MIN = 0.0           # 最小海拔高度
    H_MAX = 2000.0        # 最大海拔高度
    H_STEP = 200.0        # 海拔高度步长
    
    R_TARGETS = np.arange(R_MIN, R_MAX + R_STEP, R_STEP)
    H_TARGETS = np.arange(H_MIN, H_MAX + H_STEP, H_STEP)

    if len(R_TARGETS) < 2 or len(H_TARGETS) < 2:
        print("❌ 错误: 步长设置不合理，无法生成表格。")
        return

    # 存储结果的二维数组 (存储等效射程与目标射程的差值: R_eq - R_target)
    results_table = np.zeros((len(H_TARGETS), len(R_TARGETS)))

    print("--- 🎯 正在进行 2D 海拔高度最优发射参数搜索及等效射程计算 ---")
            
    # 填充表格
    for j, R_target in enumerate(R_TARGETS):
        print(f"\n[目标射程 R={R_target:.0f}m] 正在搜索最优参数...")
        
        for i, H_target in enumerate(H_TARGETS):
            
            # 1. 搜索最优的 (V0, Theta) 组合，使其在 H_target 下射程达到 R_target
            V0_opt, Theta_opt = search_min_v0_and_theta(R_target, DT, h0=H_target)
            
            R_diff = np.nan
            if not np.isnan(V0_opt):
                # 2. 计算等效射程 (Equivalent Range)
                # 使用找到的最优 V0 和 Theta，但在海平面 (h0=0.0) 条件下进行模拟
                try:
                    R_equivalent = simulate_and_get_impact_2d(V0_opt, Theta_opt, DT, h0=0.0)
                    # 3. 计算差值 R_eq - R_target
                    R_diff = R_equivalent - R_target
                except Exception as e:
                    print(f"  警告: H={H_target}m, R={R_target}m 模拟 R_eq 失败: {e}")
                    R_diff = np.nan
            
            # 4. 存储差值
            results_table[i, j] = R_diff
            
            # 打印调试信息
            V0_str = f"{V0_opt:6.2f}" if not np.isnan(V0_opt) else " N/A "
            Theta_str = f"{Theta_opt:4.1f}" if not np.isnan(Theta_opt) else "N/A"
            R_diff_str = f"{R_diff:+6.1f}" if not np.isnan(R_diff) else " N/A "
            
            print(f"  H={H_target:4.0f}m => V₀:{V0_str}m/s, θ:{Theta_str}° | R_eq - R:{R_diff_str}m")


    # --- 结果表格输出 ---
    print("\n" + "="*105)
    print("      🎯 海拔高度对所需“弹道能量”的影响查表 (等效射程差值: $R_{\\text{eq}} - R_{\\text{target}}$) 🎯")
    print("    表格内容为: (在海平面发射等效射程 $R_{\\text{eq}}$) - (目标射程 $R_{\\text{target}}$) (单位: 米)")
    print("    [注意: $R_{\\text{eq}}$ 使用该情形下最优 $V₀/θ$ 组合，但在 $H=0$ 条件下计算]")
    print("="*105)

    # 格式化表头 (X轴：目标直线距离 R (m))
    header_cols = [f"{r:^15.0f}" for r in R_TARGETS]
    header = " | " + " | ".join(header_cols) + " |"
    print(" 海拔高度 H (m) \\ 目标距离 R (m)")
    print("-" * len(header))
    print(header)
    print("-" * len(header))
    
    # 格式化数据行 (Y轴：海拔高度 H (m))
    for i, H in enumerate(H_TARGETS):
        row = []
        for j in range(len(R_TARGETS)):
            val = results_table[i, j]
            if np.isnan(val):
                row.append(f"{'N/A':^15}")
            else:
                # 差值保留一位小数，并显示正负号
                row.append(f"{val:^+15.1f}")
        
        print(f" | {H:^10.0f} | " + " | ".join(row) + " |")

    print("="*105)
    print("💡 查表指引:")
    print("   - 正值 (+): 表示由于海拔 H 升高，空气阻力减小，实现目标 R_{target} 所需的 V₀/θ 组合在海平面 (H=0) 本可打出更远的距离 $R_{\\text{eq}}$。")
    print("   - 负值 (-): 表示实现目标 $R_{\\text{target}}$ 所需的 V₀/θ 组合在海平面 (H=0) 只能打出更近的距离 $R_{\\text{eq}}$ (这在H>0的情形下通常不发生)。")

if __name__ == "__main__":
    generate_altitude_table()